"""
Demo 03: 频域特征（FFT 幅值）
使用快速傅里叶变换提取时间序列的频域特征
"""

import numpy as np
from demo_01_raw import RawVectorizer
import matplotlib.pyplot as plt


class FFTVectorizer:
    """FFT频域特征向量化处理类"""
    
    def __init__(self, window_size=60, step_size=5, n_coeff=12):
        """
        初始化参数
        
        Args:
            window_size: 窗口大小
            step_size: 步长
            n_coeff: 保留的FFT系数数量
        """
        self.window_size = window_size
        self.step_size = step_size
        self.n_coeff = n_coeff
        self.raw_vectorizer = RawVectorizer(window_size, step_size)
        
    def extract_fft_features(self, windows, n_coeff=None):
        """
        从窗口中提取FFT频域特征
        
        Args:
            windows: 窗口矩阵 (n_windows, window_size)
            n_coeff: 保留的系数数量
            
        Returns:
            numpy.ndarray: FFT特征矩阵 (n_windows, n_coeff)
        """
        if n_coeff is None:
            n_coeff = self.n_coeff
            
        # 对每个窗口进行FFT变换
        fft_result = np.fft.rfft(windows, axis=1)
        
        # 计算幅值（忽略相位信息）
        amplitudes = np.abs(fft_result)
        
        # 丢掉直流分量（第0项），保留前n_coeff个频率分量
        fft_features = amplitudes[:, 1:n_coeff+1]
        
        return fft_features
    
    def extract_extended_fft_features(self, windows):
        """
        提取扩展的FFT特征（包括幅值和相位）
        
        Args:
            windows: 窗口矩阵
            
        Returns:
            numpy.ndarray: 扩展FFT特征矩阵
        """
        # 对每个窗口进行FFT变换
        fft_result = np.fft.rfft(windows, axis=1)
        
        # 计算幅值和相位
        amplitudes = np.abs(fft_result)[:, 1:self.n_coeff+1]
        phases = np.angle(fft_result)[:, 1:self.n_coeff+1]
        
        # 额外特征：功率谱密度
        psd = amplitudes ** 2
        
        # 频谱统计特征
        spectral_centroid = self.compute_spectral_centroid(amplitudes)
        spectral_spread = self.compute_spectral_spread(amplitudes, spectral_centroid)
        
        # 合并所有特征
        features = np.hstack([
            amplitudes,
            phases,
            psd.mean(axis=1, keepdims=True),
            spectral_centroid.reshape(-1, 1),
            spectral_spread.reshape(-1, 1)
        ])
        
        return features
    
    def compute_spectral_centroid(self, amplitudes):
        """
        计算频谱质心
        
        Args:
            amplitudes: FFT幅值矩阵
            
        Returns:
            numpy.ndarray: 频谱质心
        """
        freqs = np.arange(amplitudes.shape[1])
        centroid = np.sum(freqs * amplitudes, axis=1) / np.sum(amplitudes, axis=1)
        return centroid
    
    def compute_spectral_spread(self, amplitudes, centroid):
        """
        计算频谱扩散度
        
        Args:
            amplitudes: FFT幅值矩阵
            centroid: 频谱质心
            
        Returns:
            numpy.ndarray: 频谱扩散度
        """
        freqs = np.arange(amplitudes.shape[1])
        spread = np.sqrt(
            np.sum(((freqs - centroid.reshape(-1, 1)) ** 2) * amplitudes, axis=1) / 
            np.sum(amplitudes, axis=1)
        )
        return spread
    
    def get_frequency_bins(self):
        """
        获取频率箱
        
        Returns:
            numpy.ndarray: 频率箱数组
        """
        # 假设采样率为1（日数据）
        freqs = np.fft.rfftfreq(self.window_size, d=1)
        return freqs[1:self.n_coeff+1]
    
    def fit_transform(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        完整的FFT特征向量化流程
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            tuple: (原始窗口矩阵, FFT特征矩阵)
        """
        # 获取原始窗口
        X_raw = self.raw_vectorizer.fit_transform(ticker, start, end)
        
        # 提取FFT特征
        X_fft = self.extract_fft_features(X_raw)
        
        return X_raw, X_fft
    
    def plot_frequency_spectrum(self, window, title="Frequency Spectrum"):
        """
        绘制频谱图
        
        Args:
            window: 单个时间序列窗口
            title: 图表标题
        """
        # 计算FFT
        fft_result = np.fft.rfft(window)
        amplitudes = np.abs(fft_result)
        freqs = np.fft.rfftfreq(len(window), d=1)
        
        # 绘图
        plt.figure(figsize=(12, 4))
        
        # 时域图
        plt.subplot(1, 2, 1)
        plt.plot(window)
        plt.title("Time Domain")
        plt.xlabel("Time")
        plt.ylabel("Value")
        plt.grid(True)
        
        # 频域图
        plt.subplot(1, 2, 2)
        plt.stem(freqs[1:self.n_coeff+1], amplitudes[1:self.n_coeff+1])
        plt.title("Frequency Domain (FFT)")
        plt.xlabel("Frequency")
        plt.ylabel("Amplitude")
        plt.grid(True)
        
        plt.suptitle(title)
        plt.tight_layout()
        plt.show()


def main():
    """主函数，演示使用方法"""
    print("=" * 60)
    print("Demo 03: 频域特征（FFT）")
    print("=" * 60)
    
    # 创建FFT特征向量化器
    vectorizer = FFTVectorizer(window_size=60, step_size=5, n_coeff=12)
    
    # 执行向量化
    print("正在下载AAPL股票数据并提取FFT特征...")
    X_raw, X_fft = vectorizer.fit_transform(
        ticker="AAPL",
        start="2020-01-01",
        end="2024-12-31"
    )
    
    # 输出结果
    print(f"\n原始窗口 shape: {X_raw.shape}")
    print(f"FFT特征 shape: {X_fft.shape}")
    print(f"保留的频率分量数: {X_fft.shape[1]}")
    
    # 显示频率箱
    freq_bins = vectorizer.get_frequency_bins()
    print(f"\n频率箱（前5个）: {freq_bins[:5]}")
    
    # 显示第一个窗口的FFT特征
    print(f"\n第一个窗口的FFT幅值特征:")
    for i in range(min(5, X_fft.shape[1])):
        print(f"  频率 {freq_bins[i]:.4f}: 幅值 = {X_fft[0, i]:.4f}")
    
    # 统计FFT特征
    print(f"\nFFT特征统计:")
    print(f"  平均幅值: {X_fft.mean():.4f}")
    print(f"  最大幅值: {X_fft.max():.4f}")
    print(f"  最小幅值: {X_fft.min():.4f}")
    
    # 主频分析
    dominant_freq_idx = X_fft.mean(axis=0).argmax()
    print(f"\n主导频率: {freq_bins[dominant_freq_idx]:.4f} (第{dominant_freq_idx+1}个频率分量)")
    
    # 可选：绘制第一个窗口的频谱
    try:
        vectorizer.plot_frequency_spectrum(X_raw[0], "First Window Spectrum")
    except:
        print("\n(跳过频谱图绘制)")
    
    return X_raw, X_fft


if __name__ == "__main__":
    X_raw, X_fft = main()
